#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@author: liang kang
@contact: gangkanli1219@163.com
@desc: 
"""
import hashlib
import os

import cv2
import numpy as np
import tensorflow as tf
from dltools.data import BaseRecordGenerator
from dltools.data.dataset.cropper import VOCImageCropper
from dltools.utils.io import read_voc_xml
from object_detection.utils import dataset_util


def _single_class(labels, cls_name):
    return [0] * len(labels), ['1'] * len(cls_name)


def _create_example(objects, image, xml_path, shape, img_format='JPEG',
                    class_type=''):
    """
    创建一个 tensorflow example
    Parameters
    ----------
    objects
    image
    xml_path
    shape
    img_format
    class_type

    Returns
    -------

    """
    boxes, labels, cls_name = [], [], []
    for ob in objects:
        labels.append(ob['label'])
        boxes.append(ob['box'])
        cls_name.append(ob['name'].encode('utf8'))
    if class_type == 'single':
        labels, cls_name = _single_class(labels, cls_name)

    labels = np.asarray(labels) + 1
    labels = labels.tolist()
    boxes = np.asarray(boxes) / np.array([shape[:2] * 2])
    ymin, xmin, ymax, xmax = np.split(boxes, 4, 1)
    if isinstance(image, str):
        with tf.gfile.GFile(image, 'rb') as fid:
            encoded_jpg = fid.read()
    else:
        encoded_jpg = image.tobytes()
    key = hashlib.sha256(encoded_jpg).hexdigest()
    truncated, poses, difficult_obj = [], [], []
    example = tf.train.Example(
        features=tf.train.Features(feature={
            'image/height':
                dataset_util.int64_feature(shape[0]),
            'image/width':
                dataset_util.int64_feature(shape[1]),
            'image/filename':
                dataset_util.bytes_feature(
                    os.path.basename(xml_path).encode('utf8')),
            'image/source_id':
                dataset_util.bytes_feature(
                    os.path.basename(xml_path).encode('utf8')),
            'image/key/sha256':
                dataset_util.bytes_feature(key.encode('utf8')),
            'image/encoded':
                dataset_util.bytes_feature(encoded_jpg),
            'image/format':
                dataset_util.bytes_feature(img_format.encode('utf8')),
            'image/object/bbox/xmin':
                dataset_util.float_list_feature(xmin.tolist()),
            'image/object/bbox/xmax':
                dataset_util.float_list_feature(xmax.tolist()),
            'image/object/bbox/ymin':
                dataset_util.float_list_feature(ymin.tolist()),
            'image/object/bbox/ymax':
                dataset_util.float_list_feature(ymax.tolist()),
            'image/object/class/text':
                dataset_util.bytes_list_feature(cls_name),
            'image/object/class/label':
                dataset_util.int64_list_feature(labels),
            'image/object/difficult':
                dataset_util.int64_list_feature(difficult_obj),
            'image/object/truncated':
                dataset_util.int64_list_feature(truncated),
            'image/object/view':
                dataset_util.bytes_list_feature(poses),
        }))
    return example.SerializeToString()


def _create_base_data_example(image_path, xml_path, shape,
                              label_list, class_type):
    """
    当图像的长宽比为一个区间时，这样生成 example
    Parameters
    ----------
    image_path
    xml_path
    shape
    label_list
    class_type

    Returns
    -------

    """
    xml, flag = read_voc_xml(xml_path, shape, label_list)
    if flag:
        return None
    return [_create_example(xml['objects'],
                            image_path, xml_path, shape, class_type)]


def _create_crop_data_example(image_path, xml_path, shape, label_list,
                              logger, class_type):
    """
    当图像的长宽比超过一个区间时，这样生成 example
    Parameters
    ----------
    image_path
    xml_path
    shape
    label_list
    logger
    class_type

    Returns
    -------

    """
    if shape[0] * 0.75 > shape[1]:
        new_shape = (int(shape[1] / 0.75), shape[1])
        number = (shape[0] - new_shape[0]) // new_shape[0] + 2
        stride = (shape[0] - new_shape[0]) // number
    else:
        new_shape = (shape[0], int(shape[0] * 0.75))
        number = (shape[1] - new_shape[1]) // new_shape[1] + 2
        stride = (shape[1] - new_shape[1]) // number
    cropper = VOCImageCropper(
        image_path, xml_path, None, new_shape,
        (stride, stride), 1.0, label_list, logger)
    examples = []
    for objects, image in cropper:
        if objects is not None and image is not None:
            example = _create_example(objects, image, xml_path, shape, 'RAW',
                                      class_type)
            if example:
                examples.append(example)
    return examples


class OriginRecordGenerator(BaseRecordGenerator):

    def __init__(self, data, output, use_da, logger=None, class_type=''):
        if logger is not None:
            logger = logger.getChild('OriginRecordGenerator')
        super(OriginRecordGenerator, self).__init__(data, output, 10, logger)
        self._label_list = []
        self._use_da = use_da
        self._class_type = class_type
        for idx in range(100):
            self._label_list.append('package%d' % idx)

    def _encode_data(self):
        image_path = self._buf_data['raw']['image']
        xml_path = self._buf_data['raw']['xml']
        image = cv2.imread(image_path)
        shape = image.shape
        if 0.725 < shape[1] / shape[0] < 0.775 or self._use_da:
            self._buf_data['data'] = _create_base_data_example(
                image_path, xml_path, shape, self._label_list, self._class_type)
        else:
            self._buf_data['data'] = _create_crop_data_example(
                image_path, xml_path, shape, self._label_list, self._logger,
                self._class_type)

    def _write_data(self, writer):
        examples = self._buf_data['data']
        if examples is None:
            return
        elif isinstance(examples, list):
            if len(examples) == 0:
                return
            else:
                for example in examples:
                    writer.write(example)
        else:
            raise TypeError('Not this type of examples !')
